/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <gtest/gtest.h>


#define protected public
#define private public
#include "graph_optimizer/op_compiler/op_compiler.h"
#include "graph/utils/graph_utils.h"
#include "ops_kernel_store/sub_ops_store.h"
#include "fusion_manager/fusion_manager.h"
#include "adapter/common/op_store_adapter_manager.h"
#include "adapter/tbe_adapter/tbe_op_store_adapter.h"
#include "common/configuration.h"
#include "ops_store/sub_op_info_store.h"
#include "ops_store/ops_kernel_manager.h"
#undef private
#undef protected

using namespace testing;
using namespace fe;
using namespace ge;

using TbeOpStoreAdapterPtr = std::shared_ptr<TbeOpStoreAdapter>;

class UTEST_fusion_engine_op_compiler : public testing::Test
{
protected:
    void SetUp()
    {
        op_store_adapter_manager_ptr_ = std::make_shared<OpStoreAdapterManager>();
        TbeOpStoreAdapterPtr tbe_adapter_ptr = std::make_shared<TbeOpStoreAdapter>();
        op_store_adapter_manager_ptr_->map_all_op_store_adapter_.emplace(std::make_pair("tbe_op_adapter", tbe_adapter_ptr));
        ops_kernel_info_store_ptr_ = std::make_shared<FEOpsKernelInfoStore>(op_store_adapter_manager_ptr_);

        FEOpsStoreInfo ops_store_info;
        SubOpInfoStorePtr sub_ops_kernel_ptr = std::make_shared<SubOpInfoStore>(ops_store_info);

        OpContent op_content;
        op_content.op_type_ = "otherNode";
        map<std::string, std::string> in_info_map;
        in_info_map.emplace(std::make_pair("format", "NC1HWC0,NCHW"));
        in_info_map.emplace(std::make_pair("dtype", "float"));
        op_content.map_kernel_info_.emplace(std::make_pair("input0", in_info_map));
        op_content.map_kernel_info_.emplace(std::make_pair("input1", in_info_map));
        op_content.map_kernel_info_.emplace(std::make_pair("output0", in_info_map));
        op_content.map_kernel_info_.emplace(std::make_pair("output1", in_info_map));
        sub_ops_kernel_ptr->op_content_map_.emplace(std::make_pair("otherNode", op_content));

        OpKernelInfoPtr info_ptr = std::make_shared<OpKernelInfo>("otherNode");
        info_ptr->impl_type_ = EN_IMPL_HW_TBE;
        sub_ops_kernel_ptr->op_kernel_info_map_.emplace(std::make_pair("otherNode", info_ptr));

        SubOpsStorePtr sub_ops_store_ptr = std::make_shared<SubOpsStore>(op_store_adapter_manager_ptr_);
        sub_ops_store_ptr->format_dtype_querier_ptr_ =
                std::make_shared<FormatDtypeQuerier>(op_store_adapter_manager_ptr_);
        ops_kernel_info_store_ptr_->map_all_sub_store_info_.emplace(std::make_pair("tbe-builtin", sub_ops_store_ptr));
        OpsKernelManager::Instance(AI_CORE_NAME).sub_ops_kernel_map_["tbe-builtin"] = sub_ops_kernel_ptr;

        graph_ = CreateTestGraph();
        graph_cce_ = CreateCceGraph();
        graph_mix_ = CreateMixGraph();
    }

    void TearDown()
    {

    }

    static NodePtr CreateCceNode(string name, GeTensorDescPtr tensor_desc_ptr, ComputeGraphPtr graph)
    {
        OpDescPtr other_desc_ptr = std::make_shared<OpDesc>(name, "otherNode");
        //set OpDesc
        auto local_tensor_desc = tensor_desc_ptr->Clone();
        // add two input desc
        for (int i = 0; i < 2; ++i) {
            AttrUtils::SetStr(local_tensor_desc, "name", name + "In" + std::to_string(i));
            other_desc_ptr->AddInputDesc(local_tensor_desc);
        }
        // add two output desc
        for (int i = 0; i < 2; ++i) {
            AttrUtils::SetStr(local_tensor_desc, "name", name + "Out" + std::to_string(i));
            other_desc_ptr->AddOutputDesc(local_tensor_desc);
        }
        // add node from other_desc_ptr to graph
        // set attr
        AttrUtils::SetInt(other_desc_ptr, "T", DT_FLOAT);
        AttrUtils::SetInt(other_desc_ptr, "_fe_imply_type", EN_IMPL_HW_GENERAL_CCE);

        NodePtr node_other = graph->AddNode(other_desc_ptr);

        return node_other;
    }

    static NodePtr CreateOtherNode(string name, GeTensorDescPtr tensor_desc_ptr, ComputeGraphPtr graph)
    {
        OpDescPtr other_desc_ptr = std::make_shared<OpDesc>(name, "otherNode");
        //set OpDesc
        auto local_tensor_desc = tensor_desc_ptr->Clone();
        // add two input desc
        for (int i = 0; i < 2; ++i) {
            AttrUtils::SetStr(local_tensor_desc, "name", name + "In" + std::to_string(i));
            other_desc_ptr->AddInputDesc(local_tensor_desc);
        }
        // add two output desc
        for (int i = 0; i < 2; ++i) {
            AttrUtils::SetStr(local_tensor_desc, "name", name + "Out" + std::to_string(i));
            other_desc_ptr->AddOutputDesc(local_tensor_desc);
        }
        // add node from other_desc_ptr to graph
        // set attr
        AttrUtils::SetInt(other_desc_ptr, "T", DT_FLOAT);
        AttrUtils::SetInt(other_desc_ptr, "_fe_imply_type", EN_IMPL_HW_TBE);

        NodePtr node_other = graph->AddNode(other_desc_ptr);

        return node_other;
    }

    static ComputeGraphPtr CreateCceGraph()
    {
        ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test");
         // new a output GeTensorDesc
        GeTensorDescPtr general_ge_tensor_desc = std::make_shared<GeTensorDesc>();
        general_ge_tensor_desc->SetFormat(FORMAT_NCHW);
        general_ge_tensor_desc->SetDataType(DT_FLOAT);

        int total_node_num = 4;
        vector<NodePtr> nodes;
        for (int i = 0; i < total_node_num; ++i) {
            nodes.push_back(CreateCceNode("test/other" + std::to_string(i), general_ge_tensor_desc, graph));
        }
         /* add link of anchors */
        std::vector<OutDataAnchorPtr> srcs;
        std::vector<InDataAnchorPtr> dsts;
        for (int i = 0; i < total_node_num - 1; ++i) {
            srcs.push_back(nodes[i]->GetOutDataAnchor(0));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(0));
            srcs.push_back(nodes[i]->GetOutDataAnchor(1));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(1));
        }

        // add edges
        for (int i = 0; i < srcs.size(); ++i)
        {
            GraphUtils::AddEdge(srcs[i], dsts[i]);
        }

        return graph;
    }

    static ComputeGraphPtr CreateMixGraph()
    {
        ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test");
         // new a output GeTensorDesc
        GeTensorDescPtr general_ge_tensor_desc = std::make_shared<GeTensorDesc>();
        general_ge_tensor_desc->SetFormat(FORMAT_NCHW);
        general_ge_tensor_desc->SetDataType(DT_FLOAT);

        int total_node_num = 4;
        vector<NodePtr> nodes;
        for (int i = 0; i < 2; ++i) {
            nodes.push_back(CreateOtherNode("test/other" + std::to_string(i), general_ge_tensor_desc, graph));
        }
        for (int i = 2; i < total_node_num; ++i) {
            nodes.push_back(CreateCceNode("test/other" + std::to_string(i), general_ge_tensor_desc, graph));
        }
         /* add link of anchors */
        std::vector<OutDataAnchorPtr> srcs;
        std::vector<InDataAnchorPtr> dsts;
        for (int i = 0; i < total_node_num - 1; ++i) {
            srcs.push_back(nodes[i]->GetOutDataAnchor(0));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(0));
            srcs.push_back(nodes[i]->GetOutDataAnchor(1));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(1));
        }

        // add edges
        for (int i = 0; i < srcs.size(); ++i)
        {
            GraphUtils::AddEdge(srcs[i], dsts[i]);
        }

        return graph;
    }

    static ComputeGraphPtr CreateTestGraph()
    {
        ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test");
         // new a output GeTensorDesc
        GeTensorDescPtr general_ge_tensor_desc = std::make_shared<GeTensorDesc>();
        general_ge_tensor_desc->SetFormat(FORMAT_NCHW);
        general_ge_tensor_desc->SetDataType(DT_FLOAT);

        int total_node_num = 4;
        vector<NodePtr> nodes;
        for (int i = 0; i < total_node_num; ++i) {
            nodes.push_back(CreateOtherNode("test/other" + std::to_string(i), general_ge_tensor_desc, graph));
        }
         /* add link of anchors */
        std::vector<OutDataAnchorPtr> srcs;
        std::vector<InDataAnchorPtr> dsts;
        for (int i = 0; i < total_node_num - 1; ++i) {
            srcs.push_back(nodes[i]->GetOutDataAnchor(0));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(0));
            srcs.push_back(nodes[i]->GetOutDataAnchor(1));
            dsts.push_back(nodes[i + 1]->GetInDataAnchor(1));
        }

        // add edges
        for (int i = 0; i < srcs.size(); ++i)
        {
            GraphUtils::AddEdge(srcs[i], dsts[i]);
        }

        return graph;
    }


    FEOpsKernelInfoStorePtr ops_kernel_info_store_ptr_;
    OpStoreAdapterManagerPtr op_store_adapter_manager_ptr_;
    ComputeGraphPtr graph_;
    ComputeGraphPtr graph_cce_;
    ComputeGraphPtr graph_mix_;
};

TEST_F(UTEST_fusion_engine_op_compiler, save_fusion_node_found)
{
    auto op_compiler_ptr = std::make_shared<OpCompiler>("normal compiler", AI_CORE_NAME, op_store_adapter_manager_ptr_);
    auto node = graph_->GetDirectNode().at(0);
    AttrUtils::SetInt(node->GetOpDesc(), "fusion_scope", 1);

    int64_t scope_id = 1;

    ScopeNodeIdMap fusion_node_map;
    std::vector<ge::Node*> fusion_nodes;
    fusion_node_map.emplace(std::make_pair(1, fusion_nodes));

    Status status = op_compiler_ptr->AddNodeToFusionMap(*node, scope_id, fusion_node_map);

    EXPECT_EQ(fe::SUCCESS, status);
}

TEST_F(UTEST_fusion_engine_op_compiler, save_fusion_node_not_found)
{
    auto op_compiler_ptr = std::make_shared<OpCompiler>("normal compiler", AI_CORE_NAME, op_store_adapter_manager_ptr_);
    auto node = graph_->GetDirectNode().at(0);
    AttrUtils::SetInt(node->GetOpDesc(), "fusion_scope", 1);

    int64_t scope_id = 1;

    ScopeNodeIdMap fusion_node_map;
    std::vector<ge::Node*> fusion_nodes;
    fusion_node_map.emplace(std::make_pair(2, fusion_nodes));

    Status status = op_compiler_ptr->AddNodeToFusionMap(*node, scope_id, fusion_node_map);

    EXPECT_EQ(fe::SUCCESS, status);
}

TEST_F(UTEST_fusion_engine_op_compiler, setcompressweightattr_03)
{
    vector<int64_t> dim1 = {1, 64, 56, 56};
    GeShape shape1(dim1);
    GeTensorDesc tensor_desc1(shape1);
    tensor_desc1.SetOriginDataType(DT_INT8);
    tensor_desc1.SetDataType(DT_INT8);
    tensor_desc1.SetFormat(FORMAT_NCHW);
    tensor_desc1.SetOriginFormat(FORMAT_NCHW);
    tensor_desc1.SetOriginShape(shape1);

    vector<int64_t> dim2 = {256, 64, 1, 1};
    GeShape shape2(dim2);
    GeTensorDesc tensor_desc2(shape2);
    tensor_desc2.SetOriginDataType(DT_INT8);
    tensor_desc2.SetDataType(DT_INT8);
    tensor_desc2.SetFormat(FORMAT_NCHW);
    tensor_desc2.SetOriginFormat(FORMAT_NCHW);
    tensor_desc2.SetOriginShape(shape2);

    vector<int64_t> dim3 = {256};
    GeShape shape3(dim3);
    GeTensorDesc tensor_desc3(shape3);
    tensor_desc3.SetOriginDataType(DT_INT8);
    tensor_desc3.SetDataType(DT_INT8);
    tensor_desc3.SetFormat(FORMAT_NCHW);
    tensor_desc3.SetOriginFormat(FORMAT_NCHW);
    tensor_desc3.SetOriginShape(shape3);

    vector<int64_t> dim4 = {1, 256, 56, 56};
    GeShape shape4(dim4);
    GeTensorDesc tensor_desc4(shape4);
    tensor_desc4.SetOriginDataType(DT_INT8);
    tensor_desc4.SetDataType(DT_INT8);
    tensor_desc4.SetFormat(FORMAT_NCHW);
    tensor_desc4.SetOriginFormat(FORMAT_NCHW);
    tensor_desc4.SetOriginShape(shape4);

    OpDescPtr op_desc = std::make_shared<OpDesc>("conv2d2", "Conv2D");
    op_desc->AddInputDesc("input", tensor_desc1);
    op_desc->AddInputDesc("filter", tensor_desc2);
    op_desc->AddInputDesc("bias", tensor_desc3);
    op_desc->AddOutputDesc("out", tensor_desc4);

    ComputeGraphPtr graph = std::make_shared<ComputeGraph>("test");
    NodePtr node = graph->AddNode(op_desc);

    string path = "./air/test/engines/nneng/config/data/platform_config";
    string real_path = RealPath(path);
    PlatformInfoManager::Instance().platform_info_map_.clear();
    PlatformInfoManager::Instance().platform_infos_map_.clear();
    PlatformInfoManager::Instance().LoadConfigFile(real_path);
    Configuration::Instance(AI_CORE_NAME).soc_version_ = "Hi3796CV300ES";

    auto op_compiler_ptr = std::make_shared<OpCompiler>("normal compiler", AI_CORE_NAME, op_store_adapter_manager_ptr_);
    Status status = op_compiler_ptr->SetCompressWeightAttr(node);
    EXPECT_EQ(fe::SUCCESS, status);
    bool has_fe_weight_attr = ge::AttrUtils::HasAttr(op_desc, ATTR_NAME_FE_WEIGHT_COMPRESS);
    EXPECT_EQ(has_fe_weight_attr, false);
    PlatformInfoManager::Instance().platform_info_map_.clear();
    PlatformInfoManager::Instance().platform_infos_map_.clear();
}
